"""
This code is used to batch detect images in a folder.
"""
import os
import sys
import torch

import cv2
import time

from vision.ssd.config.fd_config import define_img_size

from vision.ssd.mb_tiny_RFB_fd import create_Mb_Tiny_RFB_fd, create_Mb_Tiny_RFB_fd_predictor
from vision.ssd.mb_tiny_fd import create_mb_tiny_fd, create_mb_tiny_fd_predictor


def select_device(test_device="cuda:0"):
    # 检查CUDA是否可用
    if torch.cuda.is_available():
        device = "cuda:0"
        print(f"Using CUDA device: {test_device}")
    else:
        device = "cpu"
        print("CUDA is not available. Using CPU instead.")

    return device


# 相关参数
test_device = select_device(test_device="cuda:0")
candidate_size = 1500  # 'nms candidate size'
threshold = 0.6  # 置信度得分
img_max_side = 640
label_path = "./models/voc-model-labels.txt"
model_path = "./models/pretrained/version-slim-640.pth"
net_type = os.path.basename(model_path).split("-")[1]
# print(net_type)
define_img_size(img_max_side)

class_names = [name.strip() for name in open(label_path).readlines()]

if net_type == 'slim':
    net = create_mb_tiny_fd(len(class_names), is_test=True, device=test_device)
    predictor = create_mb_tiny_fd_predictor(net, candidate_size=candidate_size, device=test_device)
elif net_type == 'RFB':
    net = create_Mb_Tiny_RFB_fd(len(class_names), is_test=True, device=test_device)
    predictor = create_Mb_Tiny_RFB_fd_predictor(net, candidate_size=candidate_size, device=test_device)
else:
    print("The net type is wrong!")


net.load(model_path)

# 打开默认摄像头（通常是索引为0）
cap = cv2.VideoCapture(0)

# 检查是否成功打开摄像头
if not cap.isOpened():
    print("无法打开摄像头")
    exit()

name = {"0": "BACKGROUND", "1": "face"}

while True:
    # 读取一帧图像
    ret, frame = cap.read()

    # 如果读取成功，ret返回TrueA
    if not ret:
        print("无法获取帧")
        break

    image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    start_time = time.time()
    boxes, labels, probs = predictor.predict(image, candidate_size / 2, threshold)
    Time = (time.time() - start_time) * 1000
    print("模型检测时间：{}ms".format(Time))

    for i in range(boxes.size(0)):
        if probs[i] >= 0.9:
            box = boxes[i, :]
            cv2.rectangle(frame, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), (0, 0, 255), 2)
            cv2.putText(frame, class_names[labels[0]], (int(box[0]), int(box[1]) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.7,
                    (0, 0, 255), 2)
    # 显示图像
    cv2.imshow('实时视频', frame)

    # 按 'q' 键退出循环
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

# 释放摄像头资源
cap.release()
# 关闭所有OpenCV窗口
cv2.destroyAllWindows()
